import multiprocessing
import os

import torch
from langchain_classic.chains.retrieval_qa.base import RetrievalQA  # 检索问答链
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import FakeEmbeddings
from langchain_community.llms import LlamaCpp
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import PromptTemplate  # 提示模板
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 设置代理以解决网络连接问题
os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"

'''
本地向量库加载数据
'''

# 1. 配置参数
# 指定本地模型路径，当前使用的是CodeQwen1.5-7B-Chat-GGUF模型
MODEL_PATH = "F:/MODEL/lmstudio-community/Qwen3-8B-GGUF/Qwen3-8B-Q8_0.gguf"  # 本地模型路径

# 指定要处理的PDF文件路径
PDF_PATH = "../example.pdf"  # PDF文档路径


# 2. 加载本地Qwen模型
def load_qwen_model():
    """
    加载本地Qwen模型函数
    使用LlamaCpp加载GGUF格式的Qwen模型，并配置相关参数以优化性能
    返回:
        LlamaCpp: 加载完成的模型实例
    """
    print("⏳ 正在加载Qwen模型...")
    try:
        # 配置模型参数以优化性能
        llm = LlamaCpp(
            model_path=MODEL_PATH,  # 指定模型文件路径
            n_ctx=32768,  # 设置上下文长度为4096
            n_threads=multiprocessing.cpu_count(),  # 使用所有CPU核心提高处理速度
            n_gpu_layers=-1 if torch.cuda.is_available() else 0,  # 根据CUDA可用性决定GPU层使用
            temperature=0.3,  # 控制生成文本的随机性，较低值使输出更确定性
            max_tokens=1024,  # 设置生成文本的最大长度
            top_p=0.9,  # 使用 nucleus sampling，只考虑累积概率达到0.9的词汇
            repeat_penalty=1.15,  # 对重复token施加惩罚，提高文本多样性
            verbose=False,  # 显示详细输出信息
            n_batch=512,  # 增加批处理大小以提高处理效率
            use_mlock=True,  # 锁定模型在RAM中以提高性能
            f16_kv=True,  # 使用半精度键值缓存降低内存使用
            logits_all=False,  # 不计算所有logits以节省计算资源
        )

        return llm
    except ImportError:
        print("❌ 未安装 llama-cpp-python，无法加载GGUF模型")
        print("请运行: pip install llama-cpp-python")
        raise
    except Exception as e:
        print(f"❌ 模型加载失败: {e}")
        raise


def load_embeddings_model():
    """
    加载嵌入模型
    
    Returns:
        embeddings: 加载的嵌入模型实例
    """
    # 根据CUDA可用性设置嵌入模型设备
    embedding_device = "cuda" if torch.cuda.is_available() else "cpu"

    # 定义多个候选嵌入模型，按优先级排列
    embedding_models = [
        ("BAAI/bge-small-zh-v1.5", {"device": embedding_device}),  # 中文优化模型
        ("sentence-transformers/all-MiniLM-L6-v2", {"device": embedding_device}),  # 英文通用模型
        ("sentence-transformers/all-mpnet-base-v2", {"device": embedding_device})  # 另一个英文模型
    ]

    embeddings = None
    # 首先尝试从本地文件加载模型，避免网络请求
    for model_name, model_kwargs in embedding_models:
        try:
            print(f"尝试从本地加载嵌入模型: {model_name}")
            # 将local_files_only参数移到model_kwargs中
            local_model_kwargs = model_kwargs.copy()
            local_model_kwargs['local_files_only'] = True
            embeddings = HuggingFaceEmbeddings(
                model_name=model_name,
                model_kwargs=local_model_kwargs,
                encode_kwargs={'normalize_embeddings': True}  # 归一化嵌入向量
            )
            print(f"✅ 成功从本地加载嵌入模型: {model_name}")
            break
        except Exception as e:
            print(f"❌ 从本地加载 {model_name} 失败: {str(e)[:100]}...")
            continue

    # 如果本地加载失败，尝试从网络加载
    if embeddings is None:
        print("💡 尝试从网络加载嵌入模型...")
        for model_name, model_kwargs in embedding_models:
            try:
                print(f"尝试从网络加载嵌入模型: {model_name}")
                embeddings = HuggingFaceEmbeddings(
                    model_name=model_name,
                    model_kwargs=model_kwargs,
                    encode_kwargs={'normalize_embeddings': True}
                )
                print(f"✅ 成功从网络加载嵌入模型: {model_name}")
                break
            except Exception as e:
                print(f"❌ 从网络加载 {model_name} 失败: {str(e)[:100]}...")
                continue

    # 如果所有模型都加载失败，使用FakeEmbeddings作为后备方案
    if embeddings is None:
        print("⚠️  所有嵌入模型加载失败，使用基础的文本嵌入方法...")
        embeddings = FakeEmbeddings(size=128)

    return embeddings


def create_qa_system(llm):
    """
    """
    # 检查是否已存在向量数据库
    if os.path.exists("../faiss_index"):
        print("💾 检测到已存在的向量数据库，正在加载...")
        # 加载嵌入模型
        embeddings = load_embeddings_model()
        # 加载已存在的向量数据库
        vector_db = FAISS.load_local("../faiss_index", embeddings, allow_dangerous_deserialization=True)
    else:
        # 加载和分割PDF文档
        print("📄 正在加载和分割PDF文档...")
        loader = PyPDFLoader(PDF_PATH)  # 创建PDF加载器实例
        # 配置文本分割器参数
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,  # 每个文本块的最大长度
            chunk_overlap=200,  # 相邻文本块之间的重叠长度
            separators=["\n\n", "\n", "。", "！", "？"]  # 定义分割符优先级
        )
        # 加载并分割文档
        docs = text_splitter.split_documents(loader.load())

        # 加载嵌入模型
        embeddings = load_embeddings_model()

        # 使用文档和嵌入模型创建向量数据库
        print("💾 正在创建向量数据库...")
        vector_db = FAISS.from_documents(docs, embeddings)

        # 保存向量数据库到本地
        vector_db.save_local("faiss_index")
        print("✅ 向量数据库已保存到本地")

    """
    一个完整的提示模板具有以下特点：
    1 角色设定：明确定义模型为"专业的财务分析师"，有助于提升回答的专业性
    2 任务描述：清晰说明模型需要完成的任务
    3 思维链指导：提供思考步骤，引导模型系统性地处理问题
    4 详细约束：
        语言要求（简体中文）
        准确性要求（基于上下文，不添加额外信息）
        格式要求（专业、简洁，控制在3句话内）
        特殊处理（数字/数据处理方式）
    5 边界处理：明确定义找不到相关信息时的处理方式
    """
    template = """
    你是一个专业的财务分析师，专门负责解读和分析公司的财务报告。你的任务是根据提供的上下文准确回答用户的问题。

    请仔细阅读以下上下文信息：
    {context}

    用户问题：{question}

    请按照以下步骤思考并回答：

    1. 首先在上下文中寻找与问题直接相关的信息
    2. 如果找到了相关信息，请基于这些信息给出准确、简洁的回答
    3. 如果没有找到相关信息，请回答"未找到相关信息"
    4. 如果问题涉及数字或数据，请确保准确引用，不要进行计算或推测

    回答要求：
    - 使用简体中文回答
    - 保持回答专业、准确、简洁
    - 如果涉及金额，请保留与原文一致的单位（如万元、亿元等）
    - 回答长度控制在3句话以内
    - 不要添加任何上下文中没有的额外信息
    - 不要解释你的思考过程，直接给出答案
    - 只输出一次答案，不要重复

    你的回答：
    """

    # 创建提示模板实例
    prompt = PromptTemplate(template=template, input_variables=["context", "question"])

    # 创建检索问答链
    print("🔗 正在创建问答链...")
    return RetrievalQA.from_chain_type(
        llm=llm,  # 使用已加载的语言模型
        chain_type="stuff",  # 使用"stuff"链类型，将所有相关文档一次性传递给模型
        retriever=vector_db.as_retriever(search_kwargs={"k": 4}),  # 检索器配置，返回最相关的4个文档
        return_source_documents=True,  # 返回源文档信息
        chain_type_kwargs={"prompt": prompt}  # 传递自定义提示模板
    )


# 4. 主函数
def main():
    """
    主函数，程序入口点
    负责加载模型、创建问答系统并执行示例问答
    """
    # 加载Qwen模型
    qwen_llm = load_qwen_model()
    print("✅ Qwen模型加载完成")

    # 创建问答系统
    qa_chain = create_qa_system(qwen_llm)

    # 定义示例问题列表
    questions = [
        "2024年营业收入",
        "2024年净利润"
    ]

    # 遍历问题并获取答案
    for question in questions:
        print(f"\n❓ 问题: {question}")
        # 调用问答链获取答案
        result = qa_chain.invoke({"query": question})
        # 清理重复的答案内容
        answer = result['result'].strip()
        # 移除重复的说明内容
        if "### 说明：" in answer:
            answer = answer.split("### 说明：")[0]
        # 移除重复的答案行
        lines = answer.split('\n')
        unique_lines = []
        for line in lines:
            line_clean = line.strip()
            if line_clean and line_clean not in unique_lines:
                unique_lines.append(line_clean)
        # 限制输出行数，避免重复
        final_answer = '\n'.join(unique_lines[:5])  # 最多显示5行
        print(f"💡 答案: {final_answer}")
        # 显示答案来源页码
        if result['source_documents']:
            print(f"📌 来源页码: {result['source_documents'][0].metadata['page'] + 1}")


# 程序入口点
if __name__ == "__main__":
    main()
